{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from io import StringIO\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ABM_DATA = np.load(\"./ABM_DATA.npy\")\n",
    "abm = torch.tensor(ABM_DATA).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = os.popen(\n",
    "    \"./Simulation --enable-ABM=true --network=BA --node=1000 --m0=5 --m=5 --dynamic=SIR --alpha=0.04 --beta=0.12\", 'r').readlines()[1]\n",
    "dataset = torch.tensor(np.delete(np.array(pd.read_csv(\n",
    "    StringIO(temp), header=None).fillna(0)).flatten(), [-1])).cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(23, device='cuda:0')\n",
      "tensor([[  0.,   2.,   3.,  ..., 214., 245., 280.],\n",
      "        [  0.,   2.,   3.,  ..., 214., 245., 280.],\n",
      "        [  0.,   2.,   3.,  ..., 214., 245., 280.],\n",
      "        ...,\n",
      "        [  0.,   2.,   3.,  ..., 214., 245., 280.],\n",
      "        [  0.,   2.,   3.,  ..., 214., 245., 280.],\n",
      "        [  0.,   2.,   3.,  ..., 214., 245., 280.]], device='cuda:0',\n",
      "       dtype=torch.float64)\n",
      "tensor([[  0.,   2.,   2.,  ...,   1.,   0.,   0.],\n",
      "        [  0.,   2.,   2.,  ...,   0.,   0.,   0.],\n",
      "        [  0.,   2.,   0.,  ...,   0.,   0.,   0.],\n",
      "        ...,\n",
      "        [  0.,   2.,   5.,  ..., 606., 583., 616.],\n",
      "        [  0.,   2.,   5.,  ..., 598., 575., 608.],\n",
      "        [  0.,   2.,   6.,  ..., 614., 582., 600.]], device='cuda:0',\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "index = dataset.argmax()\n",
    "print(index)\n",
    "train_set = torch.vstack([dataset[0:index - 2] for i in range(12096)])\n",
    "comp = abm[:, 0:index - 2]\n",
    "print(train_set)\n",
    "print(comp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([  0.,   2.,   3.,   3.,   3.,   3.,   4.,   5.,   5.,   7.,  15.,  29.,\n",
      "         43.,  55.,  71., 100., 137., 171., 214., 245., 280., 313., 326., 341.,\n",
      "        336., 322., 300., 281., 260., 229., 216., 192., 178., 163., 152., 140.,\n",
      "        125., 118., 109., 101.,  93.,  85.,  74.,  65.,  61.,  58.,  57.,  52.,\n",
      "         45.,  42.,  38.,  34.,  31.,  27.,  22.,  19.,  18.,  17.,  16.,  15.,\n",
      "         12.,  10.,   8.,   7.,   7.,   4.,   3.,   3.,   2.,   2.,   2.,   2.,\n",
      "          1.,   1.], device='cuda:0', dtype=torch.float64)\n",
      "tensor([  2.,   3.,   3.,   3.,   4.,   3.,   5.,   6.,  13.,  16.,  20.,  37.,\n",
      "         52.,  70.,  97., 129., 165., 215., 249., 280., 285., 295., 315., 310.,\n",
      "        298., 279., 287., 280., 273., 283., 301., 317., 304., 308., 318., 331.,\n",
      "        324., 316., 317., 326., 335., 324., 324., 325., 336., 341., 317., 306.,\n",
      "        321., 312., 333., 342., 338., 304., 306., 316., 333., 323., 299., 302.,\n",
      "        311., 288., 294., 325., 315., 320., 327., 325., 325., 301., 295., 280.,\n",
      "        283., 278., 278., 281., 286., 291., 289., 295., 302., 296., 298., 281.,\n",
      "        311., 323., 300., 289., 315., 324., 316., 310., 331., 321., 350., 351.,\n",
      "        330., 308., 315.], device='cuda:0', dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "prediction_index = torch.sum((comp - train_set).pow(2), dim=1).argmin()\n",
    "print(dataset)\n",
    "print(abm[prediction_index, 1:100])\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ab6440cfd5cf3d7f24334eb4f6b5d3da51e47c12602bc748820dd0f1504202af"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('Pytorch')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
